"""
Proxy management for FastMCP
"""

from fastmcp import FastMCP
from fastmcp.server.proxy import ProxyClient
from typing import Dict, List, Optional, Any
import json
import asyncio
import os
from dataclasses import dataclass, asdict
from pathlib import Path


@dataclass
class ProxyConfig:
    """Configuration for a single proxy."""

    id: str
    name: str
    description: str
    target_server: Dict[str, Any]
    enabled: bool
    mount_path: str


class ProxyManager:
    """Manages MCP proxy servers dynamically."""

    def __init__(self, mcp: FastMCP, config_file: str = "proxy_conf.json"):
        self.mcp = mcp
        self.config_file = config_file
        self.proxies: Dict[str, FastMCP] = {}
        self.proxy_configs: Dict[str, ProxyConfig] = {}
        self._load_config()

    def _load_config(self):
        """Load proxy configurations from JSON file."""
        if os.path.exists(self.config_file):
            try:
                with open(self.config_file, "r", encoding="utf-8") as f:
                    data = json.load(f)

                for proxy_data in data.get("proxies", []):
                    config = ProxyConfig(**proxy_data)
                    self.proxy_configs[config.id] = config

            except Exception as e:
                print(f"Error loading proxy config: {e}")

    def _save_config(self):
        """Save current proxy configurations to JSON file."""
        try:
            data = {
                "proxies": [asdict(config) for config in self.proxy_configs.values()]
            }
            with open(self.config_file, "w", encoding="utf-8") as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
        except Exception as e:
            print(f"Error saving proxy config: {e}")

    async def _remove_local_overrides(self, proxy_id: str):
        """Remove local override components for a proxy using FastMCP v2 methods."""
        if proxy_id not in self.proxy_configs:
            return

        config = self.proxy_configs[proxy_id]
        prefix = config.mount_path.lstrip("/")

        try:
            # 获取主服务器的所有工具
            main_tools = self.mcp.get_tools()

            # 查找并移除以此代理前缀开始的本地覆盖工具
            tools_to_remove = []
            for tool_name in main_tools.keys():
                if prefix and tool_name.startswith(f"{prefix}_"):
                    tools_to_remove.append(tool_name)
                elif not prefix:  # 无前缀的情况需要更细致的处理
                    # 这里可以基于工具的元数据或其他标识来判断
                    pass

            # 移除本地覆盖工具
            for tool_name in tools_to_remove:
                try:
                    self.mcp.remove_tool(tool_name)
                except Exception as e:
                    print(f"Warning: Could not remove override tool {tool_name}: {e}")

            if tools_to_remove:
                print(
                    f"Info: Removed {len(tools_to_remove)} local override tools for proxy {proxy_id}"
                )

        except Exception as e:
            print(f"Warning: Error removing local overrides for {proxy_id}: {e}")

    def _create_proxy_client(self, target_server: Dict[str, Any]) -> str:
        """Create proxy client configuration string based on target server."""
        server_type = target_server.get("type", "stdio")

        if server_type == "stdio":
            command = target_server.get("command", "python")
            args = target_server.get("args", [])
            cwd = target_server.get("cwd", ".")

            # FastMCP v2 ProxyClient支持完整的命令行字符串
            # 对于STDIO传输，直接传递命令和参数的组合
            if args:
                # 检查args是否包含有效的模块或脚本
                full_command = f"{command} {' '.join(args)}"
                return full_command
            else:
                return command

        elif server_type == "http":
            url = target_server.get("url")
            if not url:
                raise ValueError("HTTP proxy requires 'url' parameter")
            return url
        else:
            raise ValueError(f"Unsupported transport type: {server_type}")

    async def start_proxy(self, proxy_id: str) -> bool:
        """Start a specific proxy by ID using FastMCP v2 methods."""
        if proxy_id in self.proxies:
            return False  # Already running

        if proxy_id not in self.proxy_configs:
            return False  # Config not found

        config = self.proxy_configs[proxy_id]

        try:
            # 1. 移除任何之前的本地覆盖组件
            await self._remove_local_overrides(proxy_id)

            # 2. Create proxy client configuration
            client_config = self._create_proxy_client(config.target_server)

            # Create proxy client
            proxy_client = ProxyClient(client_config)

            # Create proxy FastMCP instance
            proxy = FastMCP.as_proxy(proxy_client, name=config.name)

            # Mount the proxy to the main MCP instance
            # FastMCP v2: mount(server, prefix=None)
            self.mcp.mount(proxy, prefix=config.mount_path.lstrip("/"))

            # Store the proxy instance
            self.proxies[proxy_id] = proxy

            # Update config to enabled
            config.enabled = True
            self._save_config()

            print(
                f"Info: Proxy {proxy_id} started and mounted at /{config.mount_path.lstrip('/')}"
            )
            return True

        except Exception as e:
            print(f"Error starting proxy {proxy_id}: {e}")
            return False

    async def stop_proxy(self, proxy_id: str) -> bool:
        """Stop a specific proxy by ID using FastMCP v2 local override strategy."""
        if proxy_id not in self.proxies:
            return False  # Not running

        try:
            proxy = self.proxies[proxy_id]
            config = self.proxy_configs[proxy_id]

            # FastMCP v2策略：无法直接unmount，但可以通过禁用组件来实现
            # 1. 获取代理的所有工具、资源和提示
            proxy_tools = proxy.get_tools()
            proxy_resources = proxy.get_resources()
            proxy_prompts = proxy.get_prompts()

            # 2. 在主服务器上创建禁用的本地覆盖组件
            prefix = config.mount_path.lstrip("/")

            # 禁用所有代理工具
            for tool_name, tool in proxy_tools.items():
                try:
                    # 创建禁用的本地工具覆盖代理工具
                    disabled_tool = tool.copy() if hasattr(tool, "copy") else tool
                    disabled_tool.disable()
                    # 使用带前缀的名称添加到主服务器
                    local_name = f"{prefix}_{tool_name}" if prefix else tool_name
                    self.mcp.add_tool(disabled_tool)
                except Exception as tool_error:
                    print(f"Warning: Could not disable tool {tool_name}: {tool_error}")

            print(
                f"Info: Proxy {proxy_id} logically stopped - {len(proxy_tools)} tools disabled"
            )

            # Clean up proxy instance
            if hasattr(proxy, "cleanup"):
                await proxy.cleanup()

            # Remove from active proxies
            del self.proxies[proxy_id]

            # Update config to disabled
            config.enabled = False
            self._save_config()

            return True

        except Exception as e:
            print(f"Error stopping proxy {proxy_id}: {e}")
            print(
                f"Warning: Using fallback method - proxy {proxy_id} stopped but components may still be accessible"
            )

            # 清理代理实例作为后备方案
            if proxy_id in self.proxies:
                del self.proxies[proxy_id]

            config = self.proxy_configs[proxy_id]
            config.enabled = False
            self._save_config()

            return True

    async def restart_proxy(self, proxy_id: str) -> bool:
        """Restart a specific proxy by ID using FastMCP v2 methods."""
        # 停止代理（如果正在运行）
        if proxy_id in self.proxies:
            await self.stop_proxy(proxy_id)

        # 重新启动代理
        return await self.start_proxy(proxy_id)

    def add_proxy_config(self, config: ProxyConfig) -> bool:
        """Add a new proxy configuration."""
        if config.id in self.proxy_configs:
            return False  # Already exists

        self.proxy_configs[config.id] = config
        self._save_config()
        return True

    def remove_proxy_config(self, proxy_id: str) -> bool:
        """Remove a proxy configuration."""
        if proxy_id not in self.proxy_configs:
            return False

        # Stop proxy if running
        if proxy_id in self.proxies:
            asyncio.create_task(self.stop_proxy(proxy_id))

        del self.proxy_configs[proxy_id]
        self._save_config()
        return True

    def get_proxy_config(self, proxy_id: str) -> Optional[ProxyConfig]:
        """Get proxy configuration by ID."""
        return self.proxy_configs.get(proxy_id)

    def list_proxy_configs(self) -> List[ProxyConfig]:
        """List all proxy configurations."""
        return list(self.proxy_configs.values())

    def list_active_proxies(self) -> List[str]:
        """List IDs of currently active proxies."""
        return list(self.proxies.keys())

    def get_proxy_status(self, proxy_id: str) -> Dict[str, Any]:
        """Get detailed status of a proxy."""
        if proxy_id not in self.proxy_configs:
            return {"error": "Proxy not found"}

        config = self.proxy_configs[proxy_id]
        is_running = proxy_id in self.proxies

        status = {
            "id": proxy_id,
            "name": config.name,
            "description": config.description,
            "mount_path": config.mount_path,
            "enabled": config.enabled,
            "running": is_running,
            "target_server": config.target_server,
        }

        if is_running:
            proxy = self.proxies[proxy_id]
            # Add runtime information if available
            status["runtime_info"] = {
                "uptime": getattr(proxy, "uptime", "unknown"),
                "requests_handled": getattr(proxy, "requests_handled", 0),
            }

        return status

    async def start_enabled_proxies(self):
        """Start all proxies that are configured as enabled."""
        for proxy_id, config in self.proxy_configs.items():
            if config.enabled and proxy_id not in self.proxies:
                await self.start_proxy(proxy_id)

    async def stop_all_proxies(self):
        """Stop all running proxies."""
        for proxy_id in list(self.proxies.keys()):
            await self.stop_proxy(proxy_id)
